""" 
Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
SPDX-License-Identifier: MIT

基于Dolphin模型的文档解析API服务器
"""

import base64
import io
import json
import os
import tempfile
import uuid
from typing import List, Optional, Union
import traceback
import time

import uvicorn
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, HTMLResponse
from omegaconf import OmegaConf
from PIL import Image
from pydantic import BaseModel, Field

from chat import DOLPHIN
from utils.utils import *
from logging_config import setup_development_logging, get_logger, APILoggerMixin


# Pydantic models for API
class DocumentElement(BaseModel):
    """文档元素"""
    label: str = Field(..., description="元素类型 (para/tab/fig)")
    text: str = Field(..., description="元素内容文本")
    bbox: List[float] = Field(..., description="边界框坐标 [x1, y1, x2, y2]")
    reading_order: int = Field(..., description="阅读顺序")
    figure_path: Optional[str] = Field(None, description="图片路径（仅图片元素）")


class PageResult(BaseModel):
    """页面解析结果"""
    page_number: int = Field(..., description="页码")
    elements: List[DocumentElement] = Field(..., description="页面元素列表")


class ParseResponse(BaseModel):
    """解析响应"""
    success: bool = Field(..., description="是否成功")
    message: str = Field(..., description="响应消息")
    task_id: str = Field(..., description="任务ID")
    total_pages: int = Field(..., description="总页数")
    results: List[PageResult] = Field(..., description="解析结果")
    processing_time: float = Field(..., description="处理时间（秒）")


class ErrorResponse(BaseModel):
    """错误响应"""
    success: bool = Field(False, description="是否成功")
    message: str = Field(..., description="错误信息")
    error_code: str = Field(..., description="错误代码")


# 初始化FastAPI应用
app = FastAPI(
    title="Dolphin文档解析API",
    description="基于Dolphin多模态模型的文档解析服务，支持图像和PDF文档的智能解析",
    version="1.0.0",
    docs_url="/docs",
    redoc_url="/redoc"
)

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 在生产环境中应该限制具体域名
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 初始化日志系统
setup_development_logging()
logger = get_logger(__name__)

# 全局变量
dolphin_model = None
config = None
temp_dir = "temp_uploads"
results_dir = "api_results"

# 创建必要的目录
os.makedirs(temp_dir, exist_ok=True)
os.makedirs(results_dir, exist_ok=True)
logger.info(f"创建输出目录: {temp_dir}, {results_dir}")


def init_model():
    """初始化Dolphin模型"""
    global dolphin_model, config
    try:
        config_path = "./config/Dolphin.yaml"
        if not os.path.exists(config_path):
            raise FileNotFoundError(f"配置文件不存在: {config_path}")
        
        logger.info("开始加载Dolphin模型...")
        config = OmegaConf.load(config_path)
        dolphin_model = DOLPHIN(config)
        logger.info("Dolphin模型初始化成功")
        return True
    except Exception as e:
        logger.error(f"模型初始化失败: {str(e)}")
        return False


def cleanup_temp_files(file_paths: List[str]):
    """清理临时文件"""
    for file_path in file_paths:
        try:
            if os.path.exists(file_path):
                os.remove(file_path)
        except Exception as e:
            print(f"清理临时文件失败: {file_path}, 错误: {str(e)}")


async def process_uploaded_file(file: UploadFile, max_batch_size: int) -> dict:
    """处理上传的文件"""
    import time
    start_time = time.time()
    
    task_id = str(uuid.uuid4())
    temp_files = []
    
    try:
        # 验证文件类型
        allowed_extensions = {'.jpg', '.jpeg', '.png', '.pdf'}
        file_ext = os.path.splitext(file.filename)[1].lower()
        
        if file_ext not in allowed_extensions:
            raise HTTPException(
                status_code=400, 
                detail=f"不支持的文件类型: {file_ext}。支持的类型: {allowed_extensions}"
            )
        
        # 保存上传的文件
        temp_file_path = os.path.join(temp_dir, f"{task_id}_{file.filename}")
        temp_files.append(temp_file_path)
        
        with open(temp_file_path, "wb") as buffer:
            content = await file.read()
            buffer.write(content)
        
        # 创建结果保存目录
        task_result_dir = os.path.join(results_dir, task_id)
        setup_output_dirs(task_result_dir)
        
        # 处理文档
        json_path, recognition_results = process_document(
            document_path=temp_file_path,
            model=dolphin_model,
            save_dir=task_result_dir,
            max_batch_size=max_batch_size
        )
        
        processing_time = time.time() - start_time
        
        # 格式化结果
        if isinstance(recognition_results, list) and len(recognition_results) > 0 and isinstance(recognition_results[0], dict) and "page_number" in recognition_results[0]:
            # 多页PDF结果
            formatted_results = []
            for page_result in recognition_results:
                elements = []
                for element in page_result["elements"]:
                    elements.append(DocumentElement(**element))
                formatted_results.append(PageResult(
                    page_number=page_result["page_number"],
                    elements=elements
                ))
            total_pages = len(formatted_results)
        else:
            # 单页图像结果
            elements = []
            for element in recognition_results:
                elements.append(DocumentElement(**element))
            formatted_results = [PageResult(page_number=1, elements=elements)]
            total_pages = 1
        
        return {
            "success": True,
            "message": "文档解析成功",
            "task_id": task_id,
            "total_pages": total_pages,
            "results": formatted_results,
            "processing_time": round(processing_time, 2)
        }
        
    except Exception as e:
        error_msg = f"处理文档时发生错误: {str(e)}"
        print(f"错误详情: {traceback.format_exc()}")
        raise HTTPException(status_code=500, detail=error_msg)
    
    finally:
        # 清理临时文件
        cleanup_temp_files(temp_files)


@app.on_event("startup")
async def startup_event():
    """应用启动时的初始化"""
    if not init_model():
        raise RuntimeError("模型初始化失败，无法启动服务")


@app.get("/")
async def root():
    """根路径"""
    return {
        "message": "Dolphin文档解析API服务",
        "version": "1.0.0", 
        "status": "running",
        "model_loaded": dolphin_model is not None,
        "endpoints": {
            "docs": "/docs",
            "web_ui": "/web",
            "health": "/health",
            "parse": "/parse",
            "parse_base64": "/parse_base64"
        }
    }


@app.get("/health")
async def health_check():
    """健康检查"""
    return {
        "status": "healthy",
        "model_loaded": dolphin_model is not None,
        "timestamp": time.time()
    }


@app.post("/parse", response_model=ParseResponse)
async def parse_document(
    file: UploadFile = File(..., description="要解析的文档文件（支持图像和PDF）"),
    max_batch_size: int = Form(4, description="批处理大小，默认为4")
):
    """
    解析文档
    
    - **file**: 上传的文档文件（支持JPG、JPEG、PNG、PDF格式）
    - **max_batch_size**: 批处理大小，控制并行处理的元素数量
    
    返回文档的结构化解析结果，包括文本段落、表格、图片等元素。
    """
    logger.info(f"收到文档解析请求: 文件名={file.filename}, 批处理大小={max_batch_size}")
    
    if dolphin_model is None:
        logger.error("模型未加载，拒绝请求")
        raise HTTPException(status_code=503, detail="模型未加载，请稍后重试")
    
    if max_batch_size < 1 or max_batch_size > 16:
        logger.warning(f"无效的批处理大小: {max_batch_size}")
        raise HTTPException(status_code=400, detail="批处理大小必须在1-16之间")
    
    start_time = time.time()
    try:
        result = await process_uploaded_file(file, max_batch_size)
        duration = time.time() - start_time
        logger.info(f"文档解析成功: 文件名={file.filename}, 耗时={duration:.2f}s, 页数={result.get('total_pages', 0)}")
        return ParseResponse(**result)
    except Exception as e:
        duration = time.time() - start_time
        logger.error(f"文档解析失败: 文件名={file.filename}, 耗时={duration:.2f}s, 错误={str(e)}")
        raise


@app.post("/parse_base64", response_model=ParseResponse)
async def parse_document_base64(
    image_data: str = Form(..., description="Base64编码的图像数据"),
    filename: str = Form("image.jpg", description="文件名"),
    max_batch_size: int = Form(4, description="批处理大小，默认为4")
):
    """
    通过Base64编码解析图像文档
    
    - **image_data**: Base64编码的图像数据
    - **filename**: 文件名（用于识别文件类型）
    - **max_batch_size**: 批处理大小
    """
    import time
    start_time = time.time()
    
    if dolphin_model is None:
        raise HTTPException(status_code=503, detail="模型未加载，请稍后重试")
    
    if max_batch_size < 1 or max_batch_size > 16:
        raise HTTPException(status_code=400, detail="批处理大小必须在1-16之间")
    
    task_id = str(uuid.uuid4())
    temp_files = []
    
    try:
        # 验证文件扩展名
        allowed_extensions = {'.jpg', '.jpeg', '.png'}
        file_ext = os.path.splitext(filename)[1].lower()
        
        if file_ext not in allowed_extensions:
            raise HTTPException(
                status_code=400,
                detail=f"Base64模式仅支持图像格式: {allowed_extensions}"
            )
        
        # 解码Base64数据
        try:
            # 移除data:image前缀（如果存在）
            if image_data.startswith('data:image'):
                image_data = image_data.split(',')[1]
            
            image_bytes = base64.b64decode(image_data)
            image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Base64数据解码失败: {str(e)}")
        
        # 创建结果保存目录
        task_result_dir = os.path.join(results_dir, task_id)
        setup_output_dirs(task_result_dir)
        
        # 处理图像
        base_name = os.path.splitext(filename)[0]
        json_path, recognition_results = process_single_image(
            image=image,
            model=dolphin_model,
            save_dir=task_result_dir,
            image_name=base_name,
            max_batch_size=max_batch_size,
            save_individual=True
        )
        
        processing_time = time.time() - start_time
        
        # 格式化结果
        elements = []
        for element in recognition_results:
            elements.append(DocumentElement(**element))
        
        formatted_results = [PageResult(page_number=1, elements=elements)]
        
        return ParseResponse(
            success=True,
            message="图像解析成功",
            task_id=task_id,
            total_pages=1,
            results=formatted_results,
            processing_time=round(processing_time, 2)
        )
        
    except HTTPException:
        raise
    except Exception as e:
        error_msg = f"处理Base64图像时发生错误: {str(e)}"
        print(f"错误详情: {traceback.format_exc()}")
        raise HTTPException(status_code=500, detail=error_msg)
    
    finally:
        cleanup_temp_files(temp_files)


# Web界面HTML模板 (简化版本)
WEB_INTERFACE_HTML = '''
<!DOCTYPE html>
<html lang="zh-CN">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>🐬 Dolphin文档解析</title>
    <style>
        * { margin: 0; padding: 0; box-sizing: border-box; }
        body { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); min-height: 100vh; padding: 20px; }
        .container { max-width: 1200px; margin: 0 auto; background: white; border-radius: 15px; box-shadow: 0 20px 40px rgba(0,0,0,0.1); overflow: hidden; }
        .header { background: linear-gradient(45deg, #667eea, #764ba2); color: white; padding: 30px; text-align: center; }
        .header h1 { font-size: 2.5em; margin-bottom: 10px; }
        .main-content { padding: 40px; }
        .upload-section { background: #f8f9fa; border-radius: 10px; padding: 30px; margin-bottom: 30px; border: 2px dashed #ddd; text-align: center; }
        .upload-btn { background: linear-gradient(45deg, #667eea, #764ba2); color: white; border: none; padding: 15px 30px; border-radius: 25px; font-size: 1.1em; cursor: pointer; margin: 10px; }
        .file-input { display: none; }
        .loading { display: none; text-align: center; padding: 30px; }
        .spinner { border: 4px solid #f3f3f3; border-top: 4px solid #667eea; border-radius: 50%; width: 40px; height: 40px; animation: spin 1s linear infinite; margin: 0 auto 20px; }
        @keyframes spin { 0% { transform: rotate(0deg); } 100% { transform: rotate(360deg); } }
        .results { margin-top: 30px; }
        .element { margin: 15px 0; padding: 15px; border-radius: 8px; border-left: 4px solid #667eea; background: #f8f9fa; }
        .error { background: #fff5f5; border: 1px solid #fed7d7; color: #c53030; padding: 15px; border-radius: 8px; margin: 20px 0; }
    </style>
</head>
<body>
    <div class="container">
        <div class="header">
            <h1>🐬 Dolphin文档解析</h1>
            <p>智能文档解析 - 支持图像和PDF格式</p>
        </div>
        <div class="main-content">
            <div class="upload-section">
                <h3>📁 选择文档文件</h3>
                <p>支持JPG、JPEG、PNG、PDF格式</p>
                <input type="file" id="fileInput" class="file-input" accept=".jpg,.jpeg,.png,.pdf">
                <button class="upload-btn" onclick="document.getElementById('fileInput').click()">📂 选择文件</button>
                <button class="upload-btn" onclick="processFile()" id="processBtn" style="display:none;">🚀 开始解析</button>
            </div>
            <div class="loading" id="loading">
                <div class="spinner"></div>
                <p>正在解析文档，请稍候...</p>
            </div>
            <div class="results" id="results"></div>
        </div>
    </div>
    <script>
        let selectedFile = null;
        document.getElementById('fileInput').addEventListener('change', (e) => {
            if (e.target.files.length > 0) {
                selectedFile = e.target.files[0];
                document.getElementById('processBtn').style.display = 'inline-block';
            }
        });
        async function processFile() {
            if (!selectedFile) return;
            document.getElementById('loading').style.display = 'block';
            document.getElementById('results').innerHTML = '';
            try {
                const formData = new FormData();
                formData.append('file', selectedFile);
                formData.append('max_batch_size', '4');
                const response = await fetch('/parse', { method: 'POST', body: formData });
                const result = await response.json();
                if (result.success) {
                    let html = '<h3>📊 解析结果</h3>';
                    result.results.forEach(page => {
                        html += `<h4>📖 第 ${page.page_number} 页</h4>`;
                        page.elements.forEach(element => {
                            const typeMap = { 'para': '📝 段落', 'tab': '📊 表格', 'fig': '🖼️ 图片' };
                            html += `<div class="element"><strong>${typeMap[element.label] || element.label}:</strong><br>${element.text}</div>`;
                        });
                    });
                    document.getElementById('results').innerHTML = html;
                } else {
                    document.getElementById('results').innerHTML = `<div class="error">❌ ${result.message}</div>`;
                }
            } catch (error) {
                document.getElementById('results').innerHTML = `<div class="error">❌ 网络错误: ${error.message}</div>`;
            } finally {
                document.getElementById('loading').style.display = 'none';
            }
        }
    </script>
</body>
</html>
'''


@app.get("/web", response_class=HTMLResponse)
async def web_interface():
    """Web界面"""
    return HTMLResponse(content=WEB_INTERFACE_HTML)


@app.exception_handler(Exception)
async def global_exception_handler(request, exc):
    """全局异常处理"""
    error_msg = f"服务器内部错误: {str(exc)}"
    logger.error(f"未处理的异常: {traceback.format_exc()}")
    
    return JSONResponse(
        status_code=500,
        content={
            "success": False,
            "message": error_msg,
            "error_code": "INTERNAL_SERVER_ERROR"
        }
    )


if __name__ == "__main__":
    import argparse
    import time
    
    parser = argparse.ArgumentParser(description="Dolphin文档解析API服务器")
    parser.add_argument("--host", default="0.0.0.0", help="服务器主机地址")
    parser.add_argument("--port", type=int, default=8000, help="服务器端口")
    parser.add_argument("--workers", type=int, default=1, help="工作进程数")
    parser.add_argument("--reload", action="store_true", help="开发模式自动重载")
    
    args = parser.parse_args()
    
    print(f"启动Dolphin文档解析API服务器...")
    print(f"服务地址: http://{args.host}:{args.port}")
    print(f"API文档: http://{args.host}:{args.port}/docs")
    
    uvicorn.run(
        "api_server:app",
        host=args.host,
        port=args.port,
        workers=args.workers,
        reload=args.reload,
        log_level="info"
    )